import torch
import numpy as np


class DepthToPointCloud:
    """
    将深度图批量反投影为点云的工具类。
    支持从相机内参 .npy 文件加载并自动应用。
    """

    def __init__(self, intrinsics_path, device='cuda'):
        """
        初始化模块，加载相机内参。

        Args:
            intrinsics_path (str): 相机内参文件路径 (.npy)
            device (str): 运行设备
        """
        self.device = device
        self.intrinsics = self._load_intrinsics(intrinsics_path).to(device)
        self.intrinsics_inv = torch.inverse(self.intrinsics).unsqueeze(0)  # [1, 3, 3]

    def _load_intrinsics(self, npy_path):
        """从 .npy 文件加载相机内参矩阵"""
        K = np.load(npy_path)
        K = torch.tensor(K, dtype=torch.float32)
        if K.shape != (3, 3):
            raise ValueError(f"相机内参矩阵形状错误，应为 [3,3]，但得到 {K.shape}")
        return K

    def __call__(self, depth_batch):
        """
        将批量深度图反投影为点云。

        Args:
            depth_batch (torch.Tensor): [B, 1, H, W] 深度图
        Returns:
            pointclouds (torch.Tensor): [B, H*W, 3] 点云
        """
        assert depth_batch.ndim == 4 and depth_batch.size(1) == 1, \
            f"depth_batch 形状必须为 [B, 1, H, W]，但得到 {depth_batch.shape}"

        B, _, H, W = depth_batch.shape
        device = depth_batch.device

        # 若 batch 大于 1，则扩展内参
        intrinsics_inv = self.intrinsics_inv.expand(B, -1, -1)

        # 构造像素坐标网格
        u = torch.arange(0, W, device=device).float()
        v = torch.arange(0, H, device=device).float()
        u_grid, v_grid = torch.meshgrid(u, v, indexing='xy')  # [W,H]
        ones = torch.ones_like(u_grid)
        pix_coords = torch.stack((u_grid, v_grid, ones), dim=0).unsqueeze(0).expand(B, -1, -1, -1)  # [B,3,H,W]

        # 扩展深度
        depth = depth_batch.expand(-1, 3, -1, -1)  # [B,3,H,W]

        # 展平
        pix_coords_flat = pix_coords.view(B, 3, -1)  # [B,3,H*W]
        depth_flat = depth.view(B, 3, -1)            # [B,3,H*W]

        # 反投影：P = K^-1 * [u,v,1]^T * depth
        cam_coords = torch.bmm(intrinsics_inv, pix_coords_flat) * depth_flat  # [B,3,H*W]

        # 重新调整形状为 [B, 3, H, W]
        cam_coords = cam_coords.view(B, 3, H, W)

        return cam_coords
